import time
from datetime import timedelta

from modelscope.trainers.hooks import Hook, Priority
from modelscope.utils.logger import get_logger

logger = get_logger()


class EnhancedTrainingHook(Hook):
    """
    增强版训练监控Hook，确保在ModelScope中正确生效
    """
    priority = Priority.VERY_HIGH  # 设置高优先级确保Hook被执行

    def __init__(self, interval=10, callback=None):
        super().__init__()
        self.interval = interval
        self.start_time = None
        self.epoch_start_time = None
        self.last_iter_time = None
        self.iter_times = []
        self.callback = callback

    def before_run(self, trainer):
        """训练开始前调用"""
        self.start_time = time.time()
        logger.info("\n🚀 训练开始!")
        logger.info(f"📊 总Epochs: {trainer._max_epochs}")
        logger.info("=" * 60)
        if self.callback:
            self.callback({
                'progress': 0,
                'status': 'training',
                'epoch': f'0/{trainer._max_epochs}',
                'iter': f'0/{len(trainer.data_loader)}'
            })

    def before_train_epoch(self, trainer):
        """每个epoch开始时调用"""
        self.epoch_start_time = time.time()
        self.last_iter_time = time.time()
        logger.info(f"\n🏁 Epoch {trainer._epoch + 1}/{trainer._max_epochs} 开始")
        if self.callback:
            self.callback({
                'progress': 0,
                'status': 'training',
                'epoch': f'{trainer._epoch + 1}/{trainer._max_epochs}',
                'iter': f'0/{len(trainer.data_loader)}'
            })

    def after_train_iter(self, trainer):
        try:
            # 1. 获取基本训练信息
            current_epoch = getattr(trainer, 'epoch', 0) + 1
            current_iter = getattr(trainer, '_iter', 0) + 1
            total_iters = len(trainer.data_loader)  # 使用确认存在的dataloader属性
            process = int(current_iter / total_iters / trainer._max_epochs * 100)
            current_iter = current_iter % total_iters

            # 2. 计算时间信息
            now = time.time()
            iter_time = now - self.last_iter_time
            self.iter_times.append(iter_time)
            avg_iter_time = sum(self.iter_times[-10:]) / len(self.iter_times[-10:])
            eta_seconds = avg_iter_time * (total_iters - current_iter)

            # 3. 获取学习率
            lr = self._get_learning_rate(trainer)

            # 4. 获取损失值
            # loss = getattr(trainer, 'metric_values', {}).get('loss', float('nan'))
            loss = trainer.train_outputs['loss'].item()

            # 5. 按间隔记录或每epoch最后一步记录
            if current_iter % self.interval == 0 or current_iter == total_iters:
                logger.info(
                    f"epoch [{current_epoch}][{current_iter}/{total_iters}]\t"
                    f"lr: {lr:.3e}, "
                    f"loss: {loss:.4f}, "
                    f"iter_time: {iter_time:.3f}s, "
                    f"eta: {str(timedelta(seconds=int(eta_seconds)))}, "
                    f"avg_iter_time: {avg_iter_time:.3f}s"
                )

                self.last_iter_time = now
            if self.callback:
                self.callback({
                    'progress': process,
                    'status': 'training',
                    'epoch': f'{current_epoch}/{trainer._max_epochs}',
                    'iter': f'{current_iter}/{len(trainer.data_loader)}'
                })
        except Exception as e:
            logger.warning(f"记录训练信息时出错: {str(e)}")

    def _get_learning_rate(self, trainer):
        """安全获取当前学习率"""
        try:
            if hasattr(trainer, 'optimizer') and trainer.optimizer is not None:
                for param_group in trainer.optimizer.param_groups:
                    return param_group['lr']
            if hasattr(trainer, 'current_lr'):
                return trainer.current_lr()
            if hasattr(trainer, 'get_lr'):
                return trainer.get_lr()[0]
        except:
            return float('nan')

    def after_train_epoch(self, trainer):
        """每个epoch结束时调用"""
        if not hasattr(trainer, 'metric_values'):
            return

        epoch_time = time.time() - self.epoch_start_time

        logger.info("\n" + "=" * 60)
        logger.info(f"✅ Epoch {trainer._epoch + 1} 完成!")
        logger.info(f"⏱️ 耗时: {str(timedelta(seconds=int(epoch_time)))}")

        # 打印评估指标（如果存在）
        if hasattr(trainer, 'evaluate_results'):
            logger.info("📊 评估指标:")
            for metric, value in trainer.evaluate_results.items():
                logger.info(f"  - {metric}: {value:.4f}")
        logger.info("=" * 60)

    def after_run(self, trainer):
        """训练结束时调用"""
        total_time = time.time() - self.start_time
        logger.info("\n🎉 训练完成!")
        logger.info(f"总训练时间: {str(timedelta(seconds=int(total_time)))}")
        logger.info("=" * 60)

        if self.callback:
            self.callback({
                'progress': 100,
                'status': 'completed',
                'epoch': f'{trainer._max_epochs}/{trainer._max_epochs}',
                'iter': f'{len(trainer.data_loader)}/{len(trainer.data_loader)}'
            })